"""
测试训练结果
"""

from parl.algorithms import DDQN

from Examples.DQNExamples.Dimension2D.Dim2DAgent import Dim2DAgent
from Examples.DQNExamples.PingPong.PingPongEnv import AIPingPongGame

# 一些超参数
from Examples.DQNExamples.PingPong2D.PingPong2DModel import PingPong2DModel
from Examples.DQNExamples.PingPong2D.PingPongEnv import PingPong2DEnv
from Examples.DQNExamples.PingPong2D.PingPongTraining import PingPongTraining


GAMMA = 0.999  # 奖励衰减率
LEARNING_RATE = 0.00001  # 学习率
SAVE_PATH = 'model/pingpong.ckpt'  # 模型保存路径

# 创建环境
env = PingPong2DEnv(AIPingPongGame())
obs = env.reset()

# 获取动作维度和状态维度
action_dim = env.action_space.n
obs_dim = obs.shape


# 创建模型
model = PingPong2DModel(act_dim=action_dim, obs_dim=obs_dim)

# 创建算法
algorithm = DDQN(model, gamma=GAMMA, lr=LEARNING_RATE)

# 创建Agent
agent = Dim2DAgent(algorithm, obs_dim, action_dim)

# 创建训练
train = PingPongTraining()

train.test(env, agent, render=True, max_steps=20000000, save_path=SAVE_PATH)
